#include "visualize_train_data.h"
#include "files_io.h"
#include <boost/filesystem.hpp>
#include <boost/algorithm/string.hpp>
#include <opencv2/opencv.hpp>

#include <iostream>
#include <fstream>
#include <cassert>
#include <vector>
#include <string>
#include <algorithm>

namespace tools
{
using std::cout;
using std::endl;

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void visualize_train_data(const std::string &data_root_dir)
{
    const int numebr_of_each_line = 10;

    std::string label_txt_dir = data_root_dir + "/train/labels/labelTxt/";
    std::string trainset_reclabel_txt_dir = data_root_dir + "/train/labels/trainset_reclabelTxt/";
    std::string image_dir = data_root_dir + "/train/images/";

    std::vector<std::string> label_txt_paths, reclabel_txt_paths, image_paths;

    get_dir_file_paths(label_txt_dir, label_txt_paths);
    get_dir_file_paths(trainset_reclabel_txt_dir, reclabel_txt_paths);
    get_dir_file_paths(image_dir, image_paths);

    std::sort(label_txt_paths.begin(), label_txt_paths.end());
    std::sort(reclabel_txt_paths.begin(), reclabel_txt_paths.end());
    std::sort(image_paths.begin(), image_paths.end());

    // cout << label_txt_paths.size() << endl;
    // cout << reclabel_txt_paths.size() << endl;
    // cout << image_paths.size() << endl;

    // read image and display the bunding box
    const cv::Scalar text_color = cv::Scalar(0, 255, 0);
    const cv::Scalar box_color = cv::Scalar(0, 0, 255);
    const int font_face = cv::FONT_HERSHEY_PLAIN;
    const float font_scale = 1.0;

    for (int i = 0; i < image_paths.size(); i++)
    {
        cv::Mat image = cv::imread(image_paths[i]);

        std::vector<vec_t> indexes_label, indexes_reclabel;
        std::vector<std::string> labels, reclabels;
        extract_info_from_txt_file(label_txt_paths[i], indexes_label, labels);
        extract_info_from_txt_file(reclabel_txt_paths[i], indexes_reclabel, reclabels);

        for (int j = 0; j < indexes_label.size(); j++)
        {
            draw_box(image, indexes_label[j], box_color);
            cv::putText(image, labels[j], cv::Point(indexes_label[j][0], indexes_label[j][1]), font_face, font_scale, text_color);
        }
        for (int j = 0; j < indexes_reclabel.size(); j++)
        {
            draw_box(image, indexes_reclabel[j], box_color);
            cv::putText(image, reclabels[j], cv::Point(indexes_reclabel[j][0], indexes_reclabel[j][1]), font_face, font_scale, text_color);
        }

        cv::imshow("LabelImage", image);
        cv::waitKey(0);
    }
}

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void extract_number_from_str(const std::string &str, std::vector<int> &nums,
                             const std::string &spacer)
{
    std::vector<std::string> vec_str;
    boost::split(vec_str, str, boost::is_any_of(spacer));

    for_each(vec_str.begin(), vec_str.end(),
             [&nums](const std::string &num_str) -> void {
                 try
                 {
                     nums.push_back(std::stoi(num_str));
                 }
                 catch (std::invalid_argument e)
                 {
                     ;
                 }
             });
}

void extract_number_from_str(const std::string &str, std::vector<float> &nums,
                             const std::string &spacer)
{
    std::vector<std::string> vec_str;
    boost::split(vec_str, str, boost::is_any_of(spacer));

    for_each(vec_str.begin(), vec_str.end(),
             [&nums](const std::string &num_str) -> void {
                 try
                 {
                     nums.push_back(std::stof(num_str));
                 }
                 catch (std::invalid_argument e)
                 {
                     ;
                 }
             });
}

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void draw_box(cv::Mat &image, const vec_t &box_indexes,
              const cv::Scalar &color, const int line_width)
{
    const int n = 8;
    std::vector<cv::Point> pts;
    for (int i = 0; i < n; i += 2)
    {
        pts.push_back(cv::Point(box_indexes[i], box_indexes[i + 1]));
    }
    cv::polylines(image, pts, true, color, line_width);
}
}